#include "forward.h"
#include <torch/types.h>
#include <ATen/cuda/CUDAContext.h>
#include "auxiliary.h"

__global__ void p2e_logproba_forward_kernel(
    const float* sxy,
    const float* oxy,
    const float* invcov,
    const float* logdet_invcov,
    const int* fids,
    float* logproba,
    int* max_index,
    const int num_faces,
    const int num_elipses,
    const int num_sample_points,
    const int num_all_points
) {
    int sample_idx = blockIdx.x * blockDim.x + threadIdx.x;
    int fid = sample_idx / num_sample_points; // sample_idx = fid * num_sample_points + sample_idx_in_face 
    if (sample_idx < num_all_points) {
        float max_logproba = -1e20; 
        float cum_logproba = 0; 
        int ellipse_count_in_face = 0;
        int max_index_in_face = -1;
        for (int ellipse_idx = 0; ellipse_idx < num_elipses; ++ellipse_idx) {
            if (fids[ellipse_idx] != fid) {
                continue;
            }
            ellipse_count_in_face += 1;
            // 1. 计算距离平方
            float d_sq_x = sxy[sample_idx * 2] - oxy[ellipse_idx * 2];
            float d_sq_y = sxy[sample_idx * 2 + 1] - oxy[ellipse_idx * 2 + 1];

            // 2. 计算指数项
            float invcov_00 = invcov[ellipse_idx * 4];
            float invcov_01 = invcov[ellipse_idx * 4 + 1];
            float invcov_10 = invcov[ellipse_idx * 4 + 2];
            float invcov_11 = invcov[ellipse_idx * 4 + 3];

            float rsq_minus_dsq = (d_sq_x * invcov_00 + d_sq_y * invcov_10) * d_sq_x + (d_sq_x * invcov_01 + d_sq_y * invcov_11) * d_sq_y;

            // 3. 计算对数概率
            float current_logproba = logdet_invcov[ellipse_idx] - rsq_minus_dsq;

            // 4. 取最大值
            if (current_logproba > max_logproba) {
                max_logproba = current_logproba;
                max_index_in_face = ellipse_idx;    
            }
        }
        logproba[sample_idx] = max_logproba; // 取最大值
        max_index[sample_idx] = max_index_in_face;
    }
}

std::tuple<torch::Tensor, torch::Tensor> p2e_logproba_forward(
    torch::Tensor sxy,
    torch::Tensor oxy,
    torch::Tensor invcov,
    torch::Tensor logdet_invcov,
    torch::Tensor fids
) {
    // 检查输入张量的类型

    TORCH_CHECK(sxy.dtype() == torch::kFloat32, "sxy must be a float32 tensor");
    TORCH_CHECK(oxy.dtype() == torch::kFloat32, "oxy must be a float32 tensor");
    TORCH_CHECK(invcov.dtype() == torch::kFloat32, "invcov must be a float32 tensor");
    TORCH_CHECK(logdet_invcov.dtype() == torch::kFloat32, "logdet_invcov must be a float32 tensor");
    TORCH_CHECK(fids.dtype() == torch::kInt32, "fids must be a int32 tensor");
    
    int num_faces = sxy.size(0);
    int num_sample_points = sxy.size(1);
    int num_all_points = num_faces * num_sample_points;
    int num_elipses = oxy.size(0);
    int num_fids = fids.size(0);
    TORCH_CHECK(num_fids == num_elipses, "num_fids must be equal to num_elipses");

    // 创建输出张量
    torch::Tensor logproba = torch::empty({num_faces, num_sample_points}, sxy.options());
    torch::Tensor max_index = torch::empty({num_faces, num_sample_points}, fids.options());

    // 计算线程块和网格大小
    int threads_per_block = getOptimalThreadsPerBlock();
    // std::cout << "threads_per_block: " << threads_per_block << std::endl;
    int blocks_per_grid = (num_all_points + threads_per_block - 1) / threads_per_block;
    
    // 调用 CUDA 核函数
    p2e_logproba_forward_kernel<<<blocks_per_grid, threads_per_block>>>(
        sxy.data_ptr<float>(),
        oxy.data_ptr<float>(),
        invcov.data_ptr<float>(),
        logdet_invcov.data_ptr<float>(),
        fids.data_ptr<int>(),
        logproba.data_ptr<float>(),
        max_index.data_ptr<int>(),
        num_faces,
        num_elipses,
        num_sample_points,
        num_all_points
    );
    
    return std::make_tuple(logproba, max_index);;
}